import torch
import torch.nn as nn


# 层标准化
class LayerNorm(nn.Module):
    # Borrowed from jekbradbury
    # https://github.com/pytorch/pytorch/issues/1959
    def __init__(self, features=768, eps=1e-6):  # 这里的 features 是一个数字，一般是 768
        super(LayerNorm, self).__init__()
        self.gamma = nn.Parameter(torch.ones(features))  # shape=torch.Size([768]) 全 1
        self.beta = nn.Parameter(torch.zeros(features))  # shape=torch.Size([768]) 全 0
        self.eps = eps  # 这一项是为了防止除数为 0
        pass

    def forward(self, x):  # x 是 (batch_size, seq_len, 768) y 也是 (batch_size, seq_len, 768)
        mean = x.mean(-1, keepdim=True)  # (batch_size, seq_len, 768) 的 mean(-1) 是 shape(batch_size, seq_len, 1)
        std = x.std(-1, keepdim=True)  # std 是 shape(batch_size, seq_len, 1)
        # gamma 全 1
        # beta 全 0
        return self.gamma * (x - mean) / (std + self.eps) + self.beta

    pass


class ConditionalLayerNorm(nn.Module):
    def __init__(self, features, condition, eps=1e-6):
        super(ConditionalLayerNorm, self).__init__()

        self.beta_ffn = nn.Linear(condition, features, bias=False)
        torch.nn.init.constant_(self.beta_ffn.weight, 0.)
        self.beta = nn.Parameter(torch.zeros(features))

        self.gamma_ffn = nn.Linear(condition, features, bias=False)
        torch.nn.init.constant_(self.gamma_ffn.weight, 0.)
        self.gamma = nn.Parameter(torch.ones(features))

        self.eps = eps
        pass

    def forward(self, x, c):
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        beta = self.beta + self.beta_ffn(c)
        gamma = self.gamma + self.gamma_ffn(c)

        return gamma * (x - mean) / (std + self.eps) + beta

    pass


# 测试一下 1e-6 是个什么数字
if __name__ == '__main__':
    a_number = 1e-6  # 其实就是 10的-6次方
    print(a_number)  # 1e-06
    print('%.10f' % a_number)  # 0.0000010000
    pass


# 测试一下：如果是一个二维 tensor 减去一个一维 tensor 会得到什么结果
if __name__ == '__main__':
    a_tensor = torch.tensor([
        [1, 3, 5, 6],
        [2, 4, 7, 9],
        [3, -1, 77, 8]
    ])
    b_tensor = torch.ones(4)
    print(a_tensor.shape)  # torch.Size([3, 4])
    print(b_tensor.shape)  # torch.Size([4])
    print(a_tensor - b_tensor)
    #  tensor([[ 0.,  2.,  4.,  5.],
    #         [ 1.,  3.,  6.,  8.],
    #         [ 2., -2., 76.,  7.]])
    pass


# 测试一下 tensor.mean() 方法
if __name__ == '__main__':
    a_tensor = torch.tensor([
        [1, 3, 5, 6],
        [2, 4, 7, 9],
        [3, -1, 77, 8]
    ], dtype=torch.float)
    print(a_tensor.mean())  # tensor(10.3333)
    print(a_tensor.mean(0))  # tensor([ 2.0000,  2.0000, 29.6667,  7.6667])
    print(a_tensor.mean(1))  # tensor([ 3.7500,  5.5000, 21.7500])
    print(a_tensor.mean(-1))  # tensor([ 3.7500,  5.5000, 21.7500])
    print(a_tensor.mean(-1, keepdim=True))
    # tensor([[ 3.7500],
    #         [ 5.5000],
    #         [21.7500]])
    pass


# 测试一下 tensor.mean() 方法
if __name__ == '__main__':
    a_tensor = torch.randn(2, 3, 4)
    print(a_tensor.mean(dim=-1).shape)  # torch.Size([2, 3])
    print(a_tensor.mean(dim=-1, keepdim=True).shape)  # torch.Size([2, 3, 1])
    pass


# 如果让一个 (2, 10, 768) 减去一个 (2, 10, 1) 的
if __name__ == '__main__':
    a_tensor = torch.randn(2, 3, 4)
    b_tensor = torch.randn(2, 3, 1)
    print(a_tensor)
    # tensor([[[-1.8419,  0.4469, -0.2664, -0.5149],
    #          [-0.4199,  1.6837,  0.1087, -0.8700],
    #          [ 0.3487,  0.1183,  0.0818,  1.4271]],
    #
    #         [[ 0.0501, -0.7189, -0.3073, -0.4954],
    #          [-1.3377, -0.1096, -0.1749, -0.1092],
    #          [-0.3207,  0.4264,  0.9683, -0.2627]]])
    print(b_tensor)
    # tensor([[[-0.6335],
    #          [-0.3828],
    #          [ 1.6615]],
    #
    #         [[ 2.1470],
    #          [-0.9518],
    #          [-2.2260]]])
    c_tensor = a_tensor - b_tensor
    print(c_tensor)
    # tensor([[[-1.2085,  1.0804,  0.3670,  0.1186],
    #          [-0.0371,  2.0665,  0.4916, -0.4872],
    #          [-1.3128, -1.5432, -1.5797, -0.2344]],
    #
    #         [[-2.0969, -2.8658, -2.4542, -2.6424],
    #          [-0.3858,  0.8422,  0.7770,  0.8427],
    #          [ 1.9053,  2.6524,  3.1943,  1.9633]]])
    print(c_tensor.shape)  # torch.Size([2, 3, 4])
    pass


# 测试一下 nn.Parameter() 效果是什么样子呢
if __name__ == '__main__':
    gamma = nn.Parameter(torch.ones(768))
    print(gamma.shape)  # torch.Size([768])
    pass
